import rclpy
from rclpy.node import Node
import numpy as np
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped, PointStamped,Pose,Point,Twist
from tf2_ros import TransformListener
from rclpy.time import Time
from nav_msgs.msg import OccupancyGrid,Odometry
import math
import heapq

class LocalPlaner(Node):

    def __init__(self):
        super().__init__('local_costmap_node')

        self.declare_parameter('robot_r', 0.2)
        self.declare_parameter('inflation_radius_value', 20)
        self.declare_parameter('goal_radius_value', 0.2)
        self.declare_parameter('plan_forward_point', 4.0)

        self.point_sub = self.create_subscription(
            Path,
            '/path',
            self.path_callback,
            10
        )

        self.odom_sub = self.create_subscription(
            Odometry,
            '/Odometry',
            self.odom_callback,
            10
        )

        self.map_sub = self.create_subscription(
            OccupancyGrid,
            '/local_costmap',
            self.map_callback,
            10
        )

        self.cmd_vel_pub = self.create_publisher(Twist, 'cmd_vel', 50)
        self.path_pub = self.create_publisher(Path, 'plan', 50)
        # 全局的路径
        self.global_path_msg=None
        # 当前位置
        self.base_link_pose=None
        self.map_msg=None
    
    # 里程计接收
    def odom_callback(self,msg):
        self.base_link_pose=msg.pose

    # 路径接收
    def path_callback(self,msg):
        self.global_path_msg=msg
              
    # 地图更新，规划入口
    def map_callback(self, msg):
        # 等待里程计和路径更新
        if (self.global_path_msg is None) or (self.base_link_pose is None):
            return
        # 更新地图消息
        self.map_msg=msg
        # 更新地图矩阵
        mapArray=np.array(msg.data).reshape((msg.info.height, msg.info.width))
        # 初始化待发布的消息类型
        cmd_vel_msg=Twist()
        path_msg = Path()
        path_msg.header.stamp=self.get_clock().now().to_msg()
        path_msg.header.frame_id='map'

        # 规划前向点
        plan_forward_point=self.get_parameter('plan_forward_point').get_parameter_value().double_value
        goal_radius_value=self.get_parameter('goal_radius_value').get_parameter_value().double_value
        goal_pose=Pose()
        # 下一个点的位置
        next_pose=None
        # 获取目标位置（最近的点）
        for i in range(len(self.global_path_msg.poses)):
            # 目标位置到自己的距离
            index=len(self.global_path_msg.poses)-i-1
            goal_pose_=self.global_path_msg.poses[index].pose
            distance_to_self_x=abs(goal_pose_.position.x-self.base_link_pose.pose.position.x)
            distance_to_self_y=abs(goal_pose_.position.y-self.base_link_pose.pose.position.y)
            
            # 如果最远距离的点小于目标阈值，发速度为0
            if index ==len(self.global_path_msg.poses)-1:
                # print(index)
                if distance_to_self_x<goal_radius_value and distance_to_self_y<goal_radius_value:
                    self.get_logger().info("goal on")
                    self.cmd_vel_pub.publish(cmd_vel_msg)
                    self.path_pub.publish(path_msg)
                    return
            else:
                dx=next_pose.position.x-goal_pose_.position.x
                dy=next_pose.position.y-goal_pose_.position.y
                goal_pose_.orientation.z=math.atan2(dy, dx)
            next_pose=goal_pose_
                
            if distance_to_self_x<plan_forward_point and distance_to_self_y<plan_forward_point:
                goal_pose=goal_pose_
                break
        # 给出起点和终点位姿
        start=[self.base_link_pose.pose.position.x,
               self.base_link_pose.pose.position.y,
               self.quaternion_to_yaw(self.base_link_pose.pose.orientation)]
        
        goal=[goal_pose.position.x,
              goal_pose.position.y,
              goal_pose.orientation.z]
        
        path = self.astar((self.world_to_map(self.base_link_pose.pose.position.x,self.base_link_pose.pose.position.y)),
                          (self.world_to_map(goal_pose.position.x,goal_pose.position.y)),mapArray)

        # 如果规划失败
        if not path:
            self.get_logger().info("can not get path")
            path=self.get_points_on_line((self.world_to_map(self.base_link_pose.pose.position.x,self.base_link_pose.pose.position.y)),
                                    (self.world_to_map(goal_pose.position.x,goal_pose.position.y)),mapArray)
            
        
        # 路径发布
        for node in path:
            pose = PoseStamped()
            pose.header.stamp=self.get_clock().now().to_msg()
            pose.header.frame_id='map'
            pose.pose.position.x,pose.pose.position.y= self.map_to_world(node[0],node[1])
            path_msg.poses.append(pose)

        self.path_pub.publish(path_msg)
        # 速度发布
        if len(path_msg.poses)>0:
            cmd_vel_msg.linear.x=(path_msg.poses[min(3,len(path_msg.poses)-1)].pose.position.x-self.base_link_pose.pose.position.x)*3
            cmd_vel_msg.linear.y=(path_msg.poses[min(3,len(path_msg.poses)-1)].pose.position.y-self.base_link_pose.pose.position.y)*3
            self.cmd_vel_pub.publish(cmd_vel_msg)
        else:
            self.cmd_vel_pub.publish(cmd_vel_msg)


    # 地图网格坐标系转换到世界坐标系
    def map_to_world(self, x, y):
        mx=x*self.map_msg.info.resolution+self.map_msg.info.origin.position.x
        my=y*self.map_msg.info.resolution+self.map_msg.info.origin.position.y
        return mx, my
    
    # 世界坐标系转换到地图坐标系
    def world_to_map(self, x, y):
        wx = x-self.map_msg.info.origin.position.x
        wy = y-self.map_msg.info.origin.position.y
        mx = int(wx / self.map_msg.info.resolution)
        my = int(wy / self.map_msg.info.resolution)
        return mx, my
    
    # 四元数转换为欧拉角
    def quaternion_to_yaw(self,orientation):
        siny_cosp = 2 * (orientation.w * orientation.z + orientation.x * orientation.y)
        cosy_cosp = 1 - 2 * (orientation.y**2 + orientation.z**2)
        yaw = math.atan2(siny_cosp, cosy_cosp)
        return yaw
    
    # 计算当前点的代价值
    def get_heuristic(self, node, end):
        x1, y1 = node
        x2, y2 = end
        return ((x1-x2)*(x1 - x2)) + ((y1-y2)*(y1 - y2))
    
    # 搜索相邻的八个点
    def get_neighbors(self, node,mapArray):
        neighbors = []
        x, y = node
        directions = [(1, 0), (0, 1), (-1, 0), (0, -1),(1, 1), (-1, 1), (1, -1), (-1, -1)]
        for direction in directions:
            new_x, new_y = x + direction[0], y + direction[1]
            if 0 <= new_x < mapArray.shape[1] and 0 <= new_y < mapArray.shape[0] and mapArray[new_y, new_x] != 100:
                neighbors.append((new_x, new_y))
        return neighbors
    
    # A*搜索
    def astar(self, start, end,mapArray):
        # 给定的起始点
        orgin_start=start

        # 路径规划的openlist和closelist
        open_list_fore = []
        closed_list_fore = set()
        g_fore = dict()

        # 前向搜索
        parent_fore=dict()
        start_x,start_y=start
        end_x,end_y=end

        # 保证开始的地方在地图里
        if start_x < 0 or start_x >= mapArray.shape[1] or start_y < 0 or start_y >= mapArray.shape[0]:
            self.get_logger().info("start position is out map")
            return None
        
        # 保证结束的地方在地图里
        if end_x < 0 or end_x >= mapArray.shape[1] or end_y < 0 or end_y >= mapArray.shape[0]:
            self.get_logger().info("end position is out map")
            return None
        
        # 脱离障碍物
        if mapArray[start_y, start_x] == 100:
            parent_fore,start=self.escape_inflation(start,mapArray)
        
        if mapArray[end_y, end_x] == 100:
            parent_back,end=self.escape_inflation(end,mapArray)
        end_x,end_y=end
        if mapArray[end_y, end_x] == 100:
            return None
        # 向队列里推送A*的起点
        heapq.heappush(open_list_fore, (0, start))
        
        # 如果队列里有数据
        while open_list_fore:
            # 向前搜
            cost_fore, current_node_fore = heapq.heappop(open_list_fore)
            # 如果当前的点在从终点向起点搜的字典里，结束，计算路径
            if current_node_fore==end:
                return self.extract_path(parent_fore,orgin_start,end)
               
            # 将当前点加入已经搜过的点集
            closed_list_fore.add(current_node_fore)
            
            for neighbor in self.get_neighbors(current_node_fore,mapArray):
                # 防止重复搜索
                if neighbor in closed_list_fore:
                    continue
                # 计算代价值
                new_cost_fore = cost_fore + self.get_heuristic(neighbor, end)
                # 如果相邻的点不在字典里或代价值小
                if neighbor not in g_fore:
                    g_fore[neighbor] = math.inf
                if new_cost_fore < g_fore[neighbor]:  
                    g_fore[neighbor] = new_cost_fore
                    parent_fore[neighbor] = current_node_fore
                    heapq.heappush(open_list_fore, (new_cost_fore, neighbor))
        return None
    
    def extract_path(self, PARENT_fore,start,end):
        """
        extract path from start and goal
        :param s_meet: meet point of bi-direction a*
        :return: path
        """

        # extract path for foreward part
        path_fore = [end]
        s = end
        robot_r_value=self.get_parameter('robot_r').get_parameter_value().double_value
        set_flag=int(robot_r_value/self.map_msg.info.resolution)
        flag=set_flag
        while True:
            try:
                s = PARENT_fore[s]
            except:
                return None
            flag+=1
            if flag>=set_flag:
                flag=0
                path_fore.append(s)
            if s == start:
                break

        return list(reversed(path_fore))
    
    def get_points_on_line(self,start, end,mapArray):
        """
        获取两点之间直线上所有的点（使用Bresenham算法）。

        :param start: 元组，表示起点坐标 (start_x, start_y)
        :param end: 元组，表示终点坐标 (end_x, end_y)
        :return: 包含路径上所有点的列表
        """
        points = []
        x0, y0 = start
        x1, y1 = end
        dx = abs(x1 - x0)
        dy = abs(y1 - y0)
        x, y = x0, y0
        sx = -1 if x0 > x1 else 1
        sy = -1 if y0 > y1 else 1
        if dx > dy:
            err = dx / 2.0
            while x != x1:
                if mapArray[y,x]==100:
                    break
                points.append((x, y))
                err -= dy
                if err < 0:
                    y += sy
                    err += dx
                x += sx
        else:
            err = dy / 2.0
            while y != y1:
                if mapArray[y,x]==100:
                    break
                points.append((x, y))
                err -= dx
                if err < 0:
                    x += sx
                    err += dy
                y += sy
        # points.append((x, y))  # 包含终点
        return points

    
    # 脱离障碍物
    def escape_inflation(self,node,mapArray):
        escape_parent=dict()
        inflation_radius_value=20
        x,y=node
        output_node=node
        min_len=math.inf
        directions = [(1, 0), (0, 1), (-1, 0), (0, -1)]

        for direction in directions:
            parent_buffer=dict()
            for i in range(1,inflation_radius_value+1):
                
                new_x = x + direction[0]*i
                new_y = y + direction[1]*i
                if new_x>mapArray.shape[1]-1 or new_y>mapArray.shape[0]-1 or new_y<0 or new_x<0:
                    continue
                parent_buffer[(new_x, new_y)]=(x + direction[0]*(i-1),y + direction[1]*(i-1))
                if mapArray[new_y, new_x] != 100:
                    if i<min_len:
                        escape_parent=parent_buffer
                        min_len=i
                        output_node=(new_x, new_y)
                    break
        return escape_parent,output_node


def main(args=None):
    rclpy.init(args=args)

    local_costmap_node = LocalPlaner()

    rclpy.spin(local_costmap_node)

    local_costmap_node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
